import numpy as np
import scipy.stats
import math


def smoothing(x, a=1e-25):
    return (1 - a) * x + a


class GMMParam(object):
    def __init__(self, dim, alpha, mu):
        self.alpha = alpha
        self.mu = np.array([mu for i in range(dim)])
        self.sigma = np.diag([1 for i in range(dim)])


class GMM(object):
    def __init__(self, num_models, dim):
        self.num_models = num_models
        self.dim = dim
        self.params = None


    def train(self, x, n_iter=50, print_params=False):
        n = len(x)
        m = self.num_models
        d = self.dim
        self.params = [GMMParam(d, 1 / m, i) for i in range(m)]

        y = np.zeros((m, n))
        for i in range(m):
            for j in range(n):
                y[i][j] = smoothing(scipy.stats.multivariate_normal.pdf(x[j], mean=self.params[i].mu, cov=self.params[i].sigma))

        for iter_count in range(n_iter):
            if print_params:
                print('Iter ' + str(iter_count + 1))
                self.print_params()

            alphas = np.array([param.alpha for param in self.params])
            y_sum = np.array([(alphas * y[:,j]).sum() for j in range(n)])

            for i in range(m):
                alphai = self.params[i].alpha
                p = alphai * y[i] / y_sum
                mu = np.average(x, weights=p, axis=0)
                self.params[i].alpha = p.sum() / n
                self.params[i].mu = mu
                self.params[i].sigma = np.average([np.dot((x[j] - mu).reshape(d, 1), (x[j] - mu).reshape(1, d)) for j in range(n)], weights=p, axis=0)

            for i in range(m):
                for j in range(n):
                    y[i][j] = smoothing(scipy.stats.multivariate_normal.pdf(x[j], mean=self.params[i].mu, cov=self.params[i].sigma))

    def print_params(self):
        for i in range(self.num_models):
            print('Gaussian ' + str(i + 1))
            print('alpha = ' + str(self.params[i].alpha))
            print('mu = ' + str(self.params[i].mu))
            print('sigma = ' + str(self.params[i].sigma))

    def probability(self, x):
        return sum([self.params[i].alpha * scipy.stats.multivariate_normal.pdf(x, mean=self.params[i].mu, cov=self.params[i].sigma) for i in range(self.num_models)])


def read_data(filename):
    data = []
    with open(filename, 'r') as f:
        for line in f:
            line = line.strip().split(',')
            line = [float(i) for i in line]
            data.append(line)
    return np.array(data)


if __name__ == '__main__':
    train1 = read_data('Train1.csv')
    gmm1 = GMM(2, 2)
    print('GMM 1')
    gmm1.train(train1, n_iter=50, print_params=False)
    gmm1.print_params()
    print()

    train2 = read_data('Train2.csv')
    gmm2 = GMM(2, 2)
    print('GMM 2')
    gmm2.train(train2, n_iter=50, print_params=False)
    gmm2.print_params()
    print()

    test1 = read_data('Test1.csv')
    num_correct = 0
    for sample in test1:
        if gmm1.probability(sample) > gmm2.probability(sample):
            num_correct += 1
    accuracy = num_correct / len(test1)
    print('Test1, num_correct = ' + str(num_correct) + ', accuracy = ' + str(accuracy))

    test2 = read_data('Test2.csv')
    num_correct = 0
    for sample in test2:
        if gmm2.probability(sample) > gmm1.probability(sample):
            num_correct += 1
    accuracy = num_correct / len(test2)
    print('Test2, num_correct = ' + str(num_correct) + ', accuracy = ' + str(accuracy))
